/*
 * This is written by Lulu Han and Yandong Zheng.
 * E-mail: locomotive_crypto@163.com
 */

//#include <iostream>
#include "she2.h"
#include <math.h>

SHE2::PrivateKey::PrivateKey(const NTL::ZZ& p, const NTL::ZZ& N, const NTL::ZZ& L,
	long km, long kr, long kq){
	
	m_p = p;
	m_N = N;
	m_L = L;
	m_km = km;
	m_kr = kr;
	m_kq = kq;
}

NTL::ZZ SHE2::PrivateKey::get_p() const{
	return m_p;
}

NTL::ZZ SHE2::PrivateKey::get_N() const{
	return m_N;
}

NTL::ZZ SHE2::PrivateKey::get_L() const{
	return m_L;
}

long SHE2::PrivateKey::get_km() const{
	return m_km;
}

long SHE2::PrivateKey::get_kr() const{
	return m_kr;
}

long SHE2::PrivateKey::get_kq() const{
	return m_kq;
}

//////////////////////////////////////////////////////////////////////////////////////////////
SHE2::PublicKey::PublicKey(const NTL::ZZ& N, const NTL::ZZ& e01, const NTL::ZZ& e02, 
	const NTL::ZZ& e1, const NTL::ZZ& en1, long km, long kr, long kl, long kp, long kq){
	
	m_N = N;
	m_e01 = e01;
	m_e02 = e02;
	m_e1 = e1;
	m_en1 = en1;
	
	m_km = km;
	m_kr = kr;
	m_kq = kq;
	m_kp = kp;
	m_kl = kl;
}

NTL::ZZ SHE2::PublicKey::get_N() const{
	return m_N;
}

NTL::ZZ SHE2::PublicKey::get_enc01() const{
	return m_e01;
}

NTL::ZZ SHE2::PublicKey::get_enc02() const{
	return m_e02;
}

NTL::ZZ SHE2::PublicKey::get_enc1() const{
	return m_e1;
}

NTL::ZZ SHE2::PublicKey::get_encn1() const{
	return m_en1;
}

long SHE2::PublicKey::get_mul_depth() const{
	
	return floor(double(m_kp)/(m_kr + m_kl)) -1;
}

long SHE2::PublicKey::get_km() const{
	return m_km;
}

long SHE2::PublicKey::get_kr() const{
	return m_kr;
}

long SHE2::PublicKey::get_kq() const{
	return m_kq;
}
	
//////////////////////////////////////////////////////////////////////////////////////////////
			
SHE2::Encryptor::Encryptor(const PublicKey& pk):m_pk(pk){
	// ...
}

void SHE2::Encryptor::encrypt(const NTL::ZZ& pt, NTL::ZZ& ct){
	
	NTL::ZZ r1, r2;
	
	//std::cout << "get_kr() = " << m_pk.get_kr() << std::endl;
	r1 = NTL::RandomLen_ZZ(m_pk.get_kr()); 
	//std::cout << "r1 = " << r1 << std::endl;
	r2 = NTL::RandomLen_ZZ(m_pk.get_kr()); 
	//std::cout << "r2 = " << r2 << std::endl;
	
	if( pt >= 0 ){
		ct = (r1*m_pk.get_enc01() + r2*m_pk.get_enc02() + pt*m_pk.get_enc1()) % m_pk.get_N();
	}
	else {
		ct = (r1*m_pk.get_enc01() + r2*m_pk.get_enc02() + NTL::abs(pt)*m_pk.get_encn1()) % m_pk.get_N();
	}
}


void SHE2::Encryptor::he_add_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_kr() - 1);
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	
	ct = (c + ct1 + ct2) % m_pk.get_N();
}

void SHE2::Encryptor::he_add_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_kr() - 1);
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	
	ct = (c + ct1 + pt1) % m_pk.get_N();
}

void SHE2::Encryptor::he_mul_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_kr() - 1);
	
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	ct = (c + ct1 * ct2) % m_pk.get_N();
}

void SHE2::Encryptor::he_mul_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct){
	NTL::ZZ r, c;
	
	r = NTL::RandomLen_ZZ(m_pk.get_kr() - 1);
	
	c = (r*m_pk.get_enc01()) % m_pk.get_N();
	ct = (c + ct1 * pt1) % m_pk.get_N();
}

//////////////////////////////////////////////////////////////////////////////////////////////
SHE2::Decryptor::Decryptor(const PrivateKey& sk):m_sk(sk){
	// ...
}

void SHE2::Decryptor::encrypt(const NTL::ZZ& pt, NTL::ZZ& ct){
	NTL::ZZ rl, rp;
	
	rl = NTL::RandomLen_ZZ(m_sk.get_kr());
	rp = NTL::RandomLen_ZZ(m_sk.get_kq());
	
	if(pt >= 0){
		ct = ((rl*m_sk.get_L() + pt)*(1 + rp*m_sk.get_p())) % m_sk.get_N();
	}
	else{
		ct = ((rl*m_sk.get_L() + (m_sk.get_L() - NTL::abs(pt)))*(1 + rp*m_sk.get_p())) % m_sk.get_N();
	}
}

void SHE2::Decryptor::decrypt(const NTL::ZZ& ct, NTL::ZZ& pt){
	NTL::ZZ t;
	
	t = (ct % m_sk.get_p()) % m_sk.get_L();
	
	if( t < (m_sk.get_L() / 2)){
		pt = t;
	}
	else{
		pt = t - m_sk.get_L();
	}
}

//////////////////////////////////////////////////////////////////////////////////////////////

void SHE2::key_gen(PrivateKey& sk, PublicKey& pk, long km,
	long kr, long kl, long kp, long kq){
	
	// Generate random values p and q
	NTL::ZZ p, q;
	p =  NTL::GenPrime_ZZ(kp);
	
	// Generate random values q1, q2, ..., qi
	long cnt_qi = ceil(double(kq)/kp);
	q = 1;
	NTL::ZZ t;
	int i = 0;
	while(i < cnt_qi){
		t = NTL::GenPrime_ZZ(kp);
		
		if( 0 == (q % t)){
			continue;
		}
		q *= t;
		++i;
	}
	
	// Generate a random value L
	NTL::ZZ L;
	L =  NTL::GenPrime_ZZ(kl);
	
	NTL::ZZ N;
	N =  p * q;
	
	// Generate the PrivateKey
	PrivateKey prikey(p, N, L, km, kr, kq);
	sk = prikey;
	
	// Generate the PublicKey
	NTL::ZZ e01, e02, e1, en1;
	NTL::ZZ rl, rp;
	
	rl = NTL::RandomLen_ZZ(kr);
	rp = NTL::RandomLen_ZZ(kq);
	e01 = ((rl*L + 0)*(1 + rp*p)) % N;
	
	rl = NTL::RandomLen_ZZ(kr);
	rp = NTL::RandomLen_ZZ(kq);
	e02 = ((rl*L + 0)*(1 + rp*p)) % N;
	
	rl = NTL::RandomLen_ZZ(kr);
	rp = NTL::RandomLen_ZZ(kq);
	e1 = ((rl*L + 1)*(1 + rp*p)) % N;
	
	rl = NTL::RandomLen_ZZ(kr);
	rp = NTL::RandomLen_ZZ(kq);
	en1 = ((rl*L + (L-1))*(1 + rp*p)) % N;
	
	PublicKey pubkey(N, e01, e02, e1, en1, km, kr, kl, kp, kq);
	pk = pubkey;
}
